import torch
from torch import nn
from .wargs import WorkerArguments

class Model(nn.Module):
    def __init__(self,opt:WorkerArguments):
        super().__init__()
        self.fea=128
        self.pred_step=opt.predict_step
        self.d_in=opt.enc_in
        self.d_out=opt.dec_in

        self.preproc=nn.Linear(opt.enc_in,self.fea)
        self.layer1=nn.Linear(self.fea*opt.input_size,opt.d_model)
        self.layer2=nn.Linear(opt.d_model,self.fea*opt.predict_step)
        self.postproc=nn.Linear(self.fea,opt.dec_in)

    def forward(self,x:torch.Tensor,x_mark,y,y_mark,pretrain):
        n,l,f=x.shape
        y1=self.preproc(x.view(-1,f))
        y2=self.layer1(y1.view(n,-1))
        y3=self.layer2(y2).view(-1,self.fea)
        return self.postproc(y3).view(n,self.pred_step,self.d_out)
